"""
File   : pydantic_class.py
Desc   : langchain 结构化输出: Pydantic class
         注意: 不支持流式输出！
Date   : 2024/12/03
Author : Tianyu Chen
https://python.langchain.com/docs/how_to/structured_output/#pydantic-class
"""

from typing import Optional
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI

# Model
llm = ChatOpenAI(model="gpt-4o-mini")

# Pydantic
class Joke(BaseModel):
    """Joke to tell user."""

    setup: str = Field(description="The setup of the joke")
    punchline: str = Field(description="The punchline to the joke")
    rating: Optional[int] = Field(
        default=None, description="How funny the joke is, from 1 to 10"
    )

# Runnable
structured_llm = llm.with_structured_output(Joke)

# invoke
a_joke = structured_llm.invoke("Tell me a joke about cats")
print(type(a_joke))  # <class '__main__.Joke'>
print(a_joke)  # setup='Why did the cat sit on the computer?' punchline='Because it wanted to keep an eye on the mouse!' rating=7
print("setup:", a_joke.setup)  # setup: Why was the cat sitting on the computer?
print("punchline:", a_joke.punchline)  # punchline: Because it wanted to keep an eye on the mouse!
print("rating:", a_joke.rating)  # rating: 7
